import os
import sys

path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../'))
sys.path.append(path)

from shell.knowledge_graph.text_classification.train import *


# 加载模型
def load_model(path, is_parent):
    # model = BiLSTM(vocab_len,
    #                embedding_size,
    #                lstm_hidden_size,
    #                lstm_layer_nums,
    #                lstm_bidirectional,
    #                drop_out,
    #                n_classes)
    if is_parent:
        model = TextCNN(5)
    else:
        model = TextCNN(15)
    if torch.cuda.is_available():
        model = model.cuda()

    model.load_state_dict(torch.load(path, map_location=torch.device('cpu')))

    return model


# 预测
def predict(string):
    index = seq2index(string)
    index = np.expand_dims(index, 0)
    x = padding_seq(index)
    parent_model = load_model('parent_model.pkl', True)
    child_model = load_model('child_model.pkl', False)
    x = torch.from_numpy(x)
    if torch.cuda.is_available():
        x = x.cuda()
    y_p = parent_model(x)
    y_c = child_model(x)
    y_p = torch.argmax(y_p[0]).cpu().data.numpy()
    y_c = torch.argmax(y_c[0]).cpu().data.numpy()
    return y_p, y_c


if __name__ == '__main__':
    a, b = predict('凌颖信息传媒有限公司投资了哪些楼盘项目')
    print('主节点：', a, '子节点：', b)
